import glob
import os.path
import shutil
import random
from pathlib import Path

import cv2

from cjc.detector.detect_head import HeadDetector
from cjc.detector.opt_tool import create_opt
from cjc.tools.data_convert.tools import create_folder_if_not_exit


class LabelSmile:
    def __init__(self, folder, train_percent=0.7):
        self.detector = None
        self.val_labels = None
        self.val_images = None
        self.train_labels = None
        self.train_images = None
        self.folder = folder
        self.train_percent = train_percent
        self.files = glob.glob(self.folder)
        self.create_target_folder()
        self.create_face_detector()

    def create_target_folder(self):
        # self.train_folder ='train')
        # self.val_folder=='val')
        folder = os.path.dirname(self.folder)
        self.train_images = f'{folder}/train/images'
        create_folder_if_not_exit(self.train_images)
        self.train_labels = f'{folder}/train/labels'
        create_folder_if_not_exit(self.train_labels)
        self.val_images = f'{folder}/val/images'
        create_folder_if_not_exit(self.val_images)
        self.val_labels = f'{folder}/val/labels'
        create_folder_if_not_exit(self.val_labels)

    def create_face_detector(self):
        opt = create_opt()
        opt.view_img = False
        self.detector = HeadDetector(opt)

    def split_data(self):
        return random.random() < self.train_percent

    def copy_files_to(self, f, labels, train=True):
        label_name = os.path.basename(f).replace('jpg', 'txt')
        if train:
            shutil.copy2(f, self.train_images)
            lf = f'{self.train_labels}/{label_name}'
        else:
            shutil.copy2(f, self.val_images)
            lf = f'{self.val_labels}/{label_name}'
        with open(lf, 'w+') as fs:
            fs.writelines('\n'.join(labels))

    @staticmethod
    def modify_cls(result, cls):
        output = []
        for r in result:
            ra = r.split(' ')
            ra[0] = str(cls)
            output.append(' '.join(ra))
        return output

    def run(self, cls=0):
        """
        参数cls 一定记得更改。
        """
        index = 0
        for f in self.files:
            cvfile = cv2.imread(f)
            result = self.detector.detect(cvfile, f)
            labeltxt = self.modify_cls(result, cls)
            if len(result) > 0:
                index += 1
                is_train = self.split_data()
                self.copy_files_to(f, labeltxt, is_train)
                print(index)


if __name__ == '__main__':
    folder = '/mnt/d/python/scrapy_face/smileSpider/images/sad/*.jpg'
    ls = LabelSmile(folder)
    # ls.clear_target()
    ls.run(2)
